# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/28 17:01 
@Author  : ZhangShenao 
@File    : calc_token_count.py 
@Desc    : 使用OpenAI的tiktoken库,计算文本的token数量

tiktoken(https://github.com/openai/tiktoken/blob/main/README.md)是OpenAI开发的一种BPE分词器

tiktoken针对不同的模型,采用了不同的编码方式,而编码方式决定了如何将文本转化成token
"""
import tiktoken


def calc_token_count(text: str, encoding_name: str = "cl100k_base"):
    """计算文本的token数量"""

    encoding = tiktoken.get_encoding(encoding_name)
    return len(encoding.encode(text))


# 根据名称加载编码
cl100k_base_encoding = tiktoken.get_encoding("cl100k_base")
print(f"encoding of cl100k_base_encoding is : {cl100k_base_encoding}")

# 根据LLM模型,获取编码
o200k_base_encoding = tiktoken.encoding_for_model("gpt-4o-mini")
print(f"encoding of gpt-4o-mini is : {o200k_base_encoding}")  # o200k_base

# 将本文编码成token,获取token id列表
tokens1 = cl100k_base_encoding.encode("tiktoken is great!")
print(tokens1)
tokens2 = o200k_base_encoding.encode("tiktoken is great!")
print(tokens2)

# 将token解码成文本
text = cl100k_base_encoding.decode(tokens=[83, 1609, 5963, 374, 2294, 0])
print(f"decoded text is: {text}")

# 尽管decode()可以应用于单个token,但对于不在utf-8边界上的token来说,解码可能会有损失或错误
# 对于单个token,decode_single_token_bytes()方法可以安全地将单个整数token转换为其表示的字节
decoded_bytes = [cl100k_base_encoding.decode_single_token_bytes(token) for token in [83, 1609, 5963, 374, 2294, 0]]
print(f"decoded bytes is: {decoded_bytes}")  # decode成byte形式

# 计算文本中的token数量
# token_count = calc_token_count(text="窗前明月光", encoding_name="o200k_base")
# print(f"token is {token_count}")
